/*
 * Software License Agreement (BSD License)
 *
 *  Point Cloud Library (PCL) - www.pointclouds.org
 *  Copyright (c) 2010, Willow Garage, Inc.
 *  Copyright (c) 2012-, Open Perception, Inc.
 *
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of the copyright holder(s) nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 * $Id$
 *
 */

#ifndef PCL_GICP_OMP_H_
#define PCL_GICP_OMP_H_

#include <pcl/registration/bfgs.h>
#include <pcl/registration/icp.h>

namespace ndt_omp
{
/** \brief GeneralizedIterativeClosestPoint is an ICP variant that implements the
 * generalized iterative closest point algorithm as described by Alex Segal et al. in
 * http://www.robots.ox.ac.uk/~avsegal/resources/papers/Generalized_ICP.pdf
 * The approach is based on using anistropic cost functions to optimize the alignment
 * after closest point assignments have been made.
 * The original code uses GSL and ANN while in ours we use an eigen mapped BFGS and
 * FLANN.
 * \author Nizar Sallem
 * \ingroup registration
 */
template <typename PointSource, typename PointTarget>
class GeneralizedIterativeClosestPoint : public pcl::IterativeClosestPoint<PointSource, PointTarget>
{
public:
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::reg_name_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::getClassName;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::indices_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::target_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::input_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::tree_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::tree_reciprocal_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::nr_iterations_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::max_iterations_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::previous_transformation_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::final_transformation_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::transformation_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::transformation_epsilon_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::converged_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::corr_dist_threshold_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::inlier_threshold_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::min_number_correspondences_;
  using pcl::IterativeClosestPoint<PointSource, PointTarget>::update_visualizer_;

  typedef pcl::PointCloud<PointSource> PointCloudSource;
  typedef typename PointCloudSource::Ptr PointCloudSourcePtr;
  typedef typename PointCloudSource::ConstPtr PointCloudSourceConstPtr;

  typedef pcl::PointCloud<PointTarget> PointCloudTarget;
  typedef typename PointCloudTarget::Ptr PointCloudTargetPtr;
  typedef typename PointCloudTarget::ConstPtr PointCloudTargetConstPtr;

  typedef pcl::PointIndices::Ptr PointIndicesPtr;
  typedef pcl::PointIndices::ConstPtr PointIndicesConstPtr;

  typedef std::vector<Eigen::Matrix3d, Eigen::aligned_allocator<Eigen::Matrix3d> > MatricesVector;
  typedef boost::shared_ptr<MatricesVector> MatricesVectorPtr;
  typedef boost::shared_ptr<const MatricesVector> MatricesVectorConstPtr;

  typedef typename pcl::Registration<PointSource, PointTarget>::KdTree InputKdTree;
  typedef typename pcl::Registration<PointSource, PointTarget>::KdTreePtr InputKdTreePtr;

  typedef boost::shared_ptr<GeneralizedIterativeClosestPoint<PointSource, PointTarget> > Ptr;
  typedef boost::shared_ptr<const GeneralizedIterativeClosestPoint<PointSource, PointTarget> >
    ConstPtr;

  typedef Eigen::Matrix<double, 6, 1> Vector6d;

  /** \brief Empty constructor. */
  GeneralizedIterativeClosestPoint()
  : k_correspondences_(20),
    gicp_epsilon_(0.001),
    rotation_epsilon_(2e-3),
    mahalanobis_(0),
    max_inner_iterations_(20)
  {
    min_number_correspondences_ = 4;
    reg_name_ = "GeneralizedIterativeClosestPoint";
    max_iterations_ = 200;
    transformation_epsilon_ = 5e-4;
    corr_dist_threshold_ = 5.;
    rigid_transformation_estimation_ = boost::bind(
      &GeneralizedIterativeClosestPoint<PointSource, PointTarget>::estimateRigidTransformationBFGS,
      this, _1, _2, _3, _4, _5);
  }

  /** \brief Provide a pointer to the input dataset
   * \param cloud the const boost shared pointer to a PointCloud message
   */
  // PCL_DEPRECATED ("[pcl::registration::GeneralizedIterativeClosestPoint::setInputCloud] setInputCloud is deprecated.
  // Please use setInputSource instead.")
  void setInputCloud(const PointCloudSourceConstPtr & cloud);

  /** \brief Provide a pointer to the input dataset
   * \param cloud the const boost shared pointer to a PointCloud message
   */
  inline void setInputSource(const PointCloudSourceConstPtr & cloud)
  {
    if (cloud->points.empty()) {
      PCL_ERROR(
        "[pcl::%s::setInputSource] Invalid or empty point cloud dataset given!\n",
        getClassName().c_str());
      return;
    }
    PointCloudSource input = *cloud;
    // Set all the point.data[3] values to 1 to aid the rigid transformation
    for (size_t i = 0; i < input.size(); ++i) input[i].data[3] = 1.0;

    pcl::IterativeClosestPoint<PointSource, PointTarget>::setInputSource(cloud);
    input_covariances_.reset();
  }

  /** \brief Provide a pointer to the covariances of the input source (if computed externally!).
   * If not set, GeneralizedIterativeClosestPoint will compute the covariances itself.
   * Make sure to set the covariances AFTER setting the input source point cloud (setting the input source point cloud
   * will reset the covariances). \param[in] target the input point cloud target
   */
  inline void setSourceCovariances(const MatricesVectorPtr & covariances)
  {
    input_covariances_ = covariances;
  }

  /** \brief Provide a pointer to the input target (e.g., the point cloud that we want to align the input source to)
   * \param[in] target the input point cloud target
   */
  inline void setInputTarget(const PointCloudTargetConstPtr & target)
  {
    pcl::IterativeClosestPoint<PointSource, PointTarget>::setInputTarget(target);
    target_covariances_.reset();
  }

  /** \brief Provide a pointer to the covariances of the input target (if computed externally!).
   * If not set, GeneralizedIterativeClosestPoint will compute the covariances itself.
   * Make sure to set the covariances AFTER setting the input source point cloud (setting the input source point cloud
   * will reset the covariances). \param[in] target the input point cloud target
   */
  inline void setTargetCovariances(const MatricesVectorPtr & covariances)
  {
    target_covariances_ = covariances;
  }

  /** \brief Estimate a rigid rotation transformation between a source and a target point cloud using an iterative
   * non-linear Levenberg-Marquardt approach.
   * \param[in] cloud_src the source point cloud dataset
   * \param[in] indices_src the vector of indices describing the points of interest in \a cloud_src
   * \param[in] cloud_tgt the target point cloud dataset
   * \param[in] indices_tgt the vector of indices describing the correspondences of the interst points from \a
   * indices_src \param[out] transformation_matrix the resultant transformation matrix
   */
  void estimateRigidTransformationBFGS(
    const PointCloudSource & cloud_src, const std::vector<int> & indices_src,
    const PointCloudTarget & cloud_tgt, const std::vector<int> & indices_tgt,
    Eigen::Matrix4f & transformation_matrix);

  /** \brief \return Mahalanobis distance matrix for the given point index */
  inline const Eigen::Matrix3d & mahalanobis(size_t index) const
  {
    assert(index < mahalanobis_.size());
    return mahalanobis_[index];
  }

  /** \brief Computes rotation matrix derivative.
   * rotation matrix is obtainded from rotation angles x[3], x[4] and x[5]
   * \return d/d_rx, d/d_ry and d/d_rz respectively in g[3], g[4] and g[5]
   * param x array representing 3D transformation
   * param R rotation matrix
   * param g gradient vector
   */
  void computeRDerivative(const Vector6d & x, const Eigen::Matrix3d & R, Vector6d & g) const;

  /** \brief Set the rotation epsilon (maximum allowable difference between two
   * consecutive rotations) in order for an optimization to be considered as having
   * converged to the final solution.
   * \param epsilon the rotation epsilon
   */
  inline void setRotationEpsilon(double epsilon) { rotation_epsilon_ = epsilon; }

  /** \brief Get the rotation epsilon (maximum allowable difference between two
   * consecutive rotations) as set by the user.
   */
  inline double getRotationEpsilon() { return (rotation_epsilon_); }

  /** \brief Set the number of neighbors used when selecting a point neighbourhood
   * to compute covariances.
   * A higher value will bring more accurate covariance matrix but will make
   * covariances computation slower.
   * \param k the number of neighbors to use when computing covariances
   */
  void setCorrespondenceRandomness(int k) { k_correspondences_ = k; }

  /** \brief Get the number of neighbors used when computing covariances as set by
   * the user
   */
  int getCorrespondenceRandomness() { return (k_correspondences_); }

  /** set maximum number of iterations at the optimization step
   * \param[in] max maximum number of iterations for the optimizer
   */
  void setMaximumOptimizerIterations(int max) { max_inner_iterations_ = max; }

  ///\return maximum number of iterations at the optimization step
  int getMaximumOptimizerIterations() { return (max_inner_iterations_); }

protected:
  /** \brief The number of neighbors used for covariances computation.
   * default: 20
   */
  int k_correspondences_;

  /** \brief The epsilon constant for gicp paper; this is NOT the convergence
   * tolerence
   * default: 0.001
   */
  double gicp_epsilon_;

  /** The epsilon constant for rotation error. (In GICP the transformation epsilon
   * is split in rotation part and translation part).
   * default: 2e-3
   */
  double rotation_epsilon_;

  /** \brief base transformation */
  Eigen::Matrix4f base_transformation_;

  /** \brief Temporary pointer to the source dataset. */
  const PointCloudSource * tmp_src_;

  /** \brief Temporary pointer to the target dataset. */
  const PointCloudTarget * tmp_tgt_;

  /** \brief Temporary pointer to the source dataset indices. */
  const std::vector<int> * tmp_idx_src_;

  /** \brief Temporary pointer to the target dataset indices. */
  const std::vector<int> * tmp_idx_tgt_;

  /** \brief Input cloud points covariances. */
  MatricesVectorPtr input_covariances_;

  /** \brief Target cloud points covariances. */
  MatricesVectorPtr target_covariances_;

  /** \brief Mahalanobis matrices holder. */
  std::vector<Eigen::Matrix3d> mahalanobis_;

  /** \brief maximum number of optimizations */
  int max_inner_iterations_;

  /** \brief compute points covariances matrices according to the K nearest
   * neighbors. K is set via setCorrespondenceRandomness() methode.
   * \param cloud pointer to point cloud
   * \param tree KD tree performer for nearest neighbors search
   * \param[out] cloud_covariances covariances matrices for each point in the cloud
   */
  template <typename PointT>
  void computeCovariances(
    typename pcl::PointCloud<PointT>::ConstPtr cloud,
    const typename pcl::search::KdTree<PointT>::ConstPtr tree, MatricesVector & cloud_covariances);

  /** \return trace of mat1^t . mat2
   * \param mat1 matrix of dimension nxm
   * \param mat2 matrix of dimension nxp
   */
  inline double matricesInnerProd(const Eigen::MatrixXd & mat1, const Eigen::MatrixXd & mat2) const
  {
    double r = 0.;
    size_t n = mat1.rows();
    // tr(mat1^t.mat2)
    for (size_t i = 0; i < n; i++)
      for (size_t j = 0; j < n; j++) r += mat1(j, i) * mat2(i, j);
    return r;
  }

  /** \brief Rigid transformation computation method  with initial guess.
   * \param output the transformed input point cloud dataset using the rigid transformation found
   * \param guess the initial guess of the transformation to compute
   */
  void computeTransformation(PointCloudSource & output, const Eigen::Matrix4f & guess);

  /** \brief Search for the closest nearest neighbor of a given point.
   * \param query the point to search a nearest neighbour for
   * \param index vector of size 1 to store the index of the nearest neighbour found
   * \param distance vector of size 1 to store the distance to nearest neighbour found
   */
  inline bool searchForNeighbors(
    const PointSource & query, std::vector<int> & index, std::vector<float> & distance) const
  {
    int k = tree_->nearestKSearch(query, 1, index, distance);
    if (k == 0) return (false);
    return (true);
  }

  /// \brief compute transformation matrix from transformation matrix
  void applyState(Eigen::Matrix4f & t, const Vector6d & x) const;

  /// \brief optimization functor structure
  struct OptimizationFunctorWithIndices : public BFGSDummyFunctor<double, 6>
  {
    OptimizationFunctorWithIndices(const GeneralizedIterativeClosestPoint * gicp)
    : BFGSDummyFunctor<double, 6>(), gicp_(gicp)
    {
    }
    double operator()(const Vector6d & x);
    void df(const Vector6d & x, Vector6d & df);
    void fdf(const Vector6d & x, double & f, Vector6d & df);

    const GeneralizedIterativeClosestPoint * gicp_;
  };

  boost::function<void(
    const pcl::PointCloud<PointSource> & cloud_src, const std::vector<int> & src_indices,
    const pcl::PointCloud<PointTarget> & cloud_tgt, const std::vector<int> & tgt_indices,
    Eigen::Matrix4f & transformation_matrix)>
    rigid_transformation_estimation_;
};
}  // namespace ndt_omp

#endif  //#ifndef PCL_GICP_H_
